# 来源 https://zhuanlan.zhihu.com/p/98644042
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import torch
from torch import nn
import sys
np.set_printoptions(suppress=True)
#数据预处理

df1 = pd.read_excel('data/trendTermData210108.xlsx')
data_raw = df1['trend1']#[0:100]    #.values.tolist()
dataset = np.array(data_raw)
print (dataset)
dataset = dataset.astype('float32').reshape(1,len(dataset))
print ('第{row}行:'.format(row=sys._getframe().f_lineno ),\
       'dataset',dataset.shape)

# # 设置不同频率的仿真信号；叠加漂移为输入，不叠加为输出
time_steps = np.linspace(0,dataset.shape[1], num=dataset.shape[1])
data=np.zeros((2,5,dataset.shape[1]))
for i in range(5):
    data[0][i]=  (500*np.sin(2 * np.pi * time_steps * 40 * (1 + 0.25 * i)) + 1.5).reshape(1,dataset.shape[1])
    # data[0][i]=(np.sin(2*np.pi*time_steps*40*(1+0.25*i))+1.5)# .reshape(1,dataset.shape[1])
    data[1][i]=data[0][i]+dataset #.reshape(1,len(dataset))

#dataset=
dataset=data.astype('float32')

max_value = np.max(dataset)
min_value = np.min(dataset)
scalar = max_value - min_value
# dataset = list(map(lambda x: x/scalar, dataset)) #将数据归一化到0~1之间
dataset = dataset/scalar
look_back=100   #10
#划分数据集
#通过前面几条的数据来预测下一条的数据，look_back设置具体的把前面几条的数据作为预测的输入data_X，而输出就是下一条data_Y
def create_dataset(dataset,look_back): # 每个的滑动窗口设置为2
    dataX, dataY=[], []
    d1,d2,d3=dataset.shape
    print('第{row}行:'.format(row=sys._getframe().f_lineno ),\
          d1,d2,d3)
    for j in range(d2):
        for i in range(d3-look_back):
            a=dataset[1][j][i:(i+look_back)]
            dataX.append(a)    # 记录窗口的值
            dataY.append(dataset[0][j][i+look_back]) # 记录除了前面两个以外的所有值作为正确的标签
    return np.array(dataX), np.array(dataY)

#创建好输入与输出 data_Y作为正确的预测值
data_X, data_Y = create_dataset(dataset,look_back)
print ('第{row}行:'.format(row=sys._getframe().f_lineno ),\
       'data_X.shape',data_X.shape)
print ('第{row}行:'.format(row=sys._getframe().f_lineno ),\
       'data_Y.shape',data_Y.shape)
#划分训练集和测试集，70%作为训练集

train_size = int(len(data_X) * 0.98)
test_size = len(data_X)-train_size

train_X = data_X[:train_size]
train_Y = data_Y[:train_size]

test_X = data_X[train_size:]
test_Y = data_Y[train_size:]

#RNN 读入的数据维度是 (seq, batch, feature)，所以要重新改变一下数据的维度，这里只有一个序列，所以 batch 是 1，而输入的 feature 就是我们希望依据的几天，这里我们定的是两个天，所以 feature 就是 2.

train_X = train_X.reshape(-1,1,look_back)
# train_X= tensor.numpy(torch.from_numpy(train_X).view(5,-1,look_back))
train_Y = train_Y.reshape(-1,1,1)
test_X = test_X.reshape(-1,1,look_back)
print ('第{row}行:'.format(row=sys._getframe().f_lineno ),\
       'train_X.shape',train_X.shape)
print ('第{row}行:'.format(row=sys._getframe().f_lineno ),\
       'train_Y.shape',train_Y.shape)

# 转化成torch 的张量
train_x = torch.from_numpy(train_X)
train_y = torch.from_numpy(train_Y)
test_x = torch.from_numpy(test_X)

#lstm 网络
class lstm_reg(nn.Module):#括号中的是python的类继承语法，父类是nn.Module类 不是参数的意思
    def __init__(self,input_size,hidden_size, output_size=1,num_layers=1): # 构造函数
        #inpu_size 是输入的样本的特征维度， hidden_size 是LSTM层的神经元个数，
        #output_size是输出的特征维度
        super(lstm_reg,self).__init__()# super用于多层继承使用，必须要有的操作

        self.rnn = nn.LSTM(input_size,hidden_size,num_layers,
                           batch_first=True
                           )# 两层LSTM网络，
        self.reg = nn.Linear(hidden_size,output_size)#把上一层总共hidden_size个的神经元的输出向量作为输入向量，然后回归到output_size维度的输出向量中

    def forward(self,x): #x是输入的数据
        x, _ = self.rnn(x)# 单个下划线表示不在意的变量，这里是LSTM网络输出的两个隐藏层状态
        s,b,h = x.shape
        x = x.contiguous().view(s*b, h)
        # x = x.reshape(s * b, h)
        x = self.reg(x)
        x = x.view(s,b,-1)#使用-1表示第三个维度自动根据原来的shape 和已经定了的s,b来确定
        return x
#我使用了GPU加速，如果不用的话需要把.cuda()给注释掉
net = lstm_reg(look_back,4)
net = net#.cuda()
criterion = nn.MSELoss()#.cuda()
optimizer = torch.optim.Adam(net.parameters(),lr=1e-2)

for e in range(50):
# 新版本中可以不使用Variable了
#     var_x = Variable(train_x).cuda()
#     var_y = Variable(train_y).cuda()

#将tensor放在GPU上面进行运算
    var_x = train_x.contiguous()#.cuda()
    var_y = train_y.contiguous()#.cuda()

    out = net(var_x)
    loss = criterion(out, var_y)

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    if (e+1)%20 == 0:
        print('Epoch: {}, Loss:{:.5f}'.format(e+1, loss.item()))
#存储训练好的模型参数
torch.save(net.state_dict(), 'example_wp_log.net_params.pkl')

net.load_state_dict(torch.load('example_wp_log.net_params.pkl'))
var_data = torch.from_numpy(test_X)#.cuda()#net在GPU上面，所以输入的测试集合也要转入到GPU上面
pred_test = net(var_data) # 测试集的预测结果
pred_test = pred_test.cpu().view(-1).data.numpy()#先转移到cpu上才能转换为numpy

#乘以原来归一化的刻度放缩回到原来的值域
origin_test_Y = test_Y*scalar
origin_pred_test = pred_test*scalar

#画图
plt.subplot(2, 1, 1)
plt.plot(origin_test_Y, 'b', label='real')
plt.plot(origin_pred_test, 'r', label='predict')
plt.legend(loc='best')
plt.subplot(2, 1, 2)
plt.plot(data[1][4][data.shape[2]-len(origin_pred_test):], label='input')
plt.legend(loc='best')
plt.show()
# plt.plot(origin_test_Y, 'b', label='real')
# plt.plot(origin_pred_test, 'r', label='prediction')
# plt.legend(loc='best')
# plt.show()
#计算MSE
#loss = criterion(out, var_y)？
true_data = origin_test_Y
true_data = np.array(true_data)
true_data = np.squeeze(true_data)  # 从二维变成一维
MSE = true_data - origin_pred_test
MSE = MSE*MSE
MSE_loss = sum(MSE)/len(MSE)
print(MSE_loss)